import os
import math
import argparse

import torch
from torch.multiprocessing import Queue, Process
import numpy as np
from tqdm import tqdm
import cv2

import network
import dataset
import misc_utils
from config import config
from visual_utils import draw_boxes

if_set_nms = True
image_list = [
    '/root/data/CrowdHuman/val/Images/282555,75b5f000bc7f1872.jpg',
    '/root/data/CrowdHuman/val/Images/283081,127530004373b993.jpg',
    '/root/data/CrowdHuman/val/Images/283647,177f00000a9633e5.jpg',
    '/root/data/CrowdHuman/val/Images/283992,76060002374133e.jpg',
    '/root/data/CrowdHuman/val/Images/284193,fe2e00004ff1627.jpg',
    '/root/data/CrowdHuman/val/Images/282555,75ba7000dfe99989.jpg',
    '/root/data/CrowdHuman/val/Images/283081,12bcb000e5a5beda.jpg',
    '/root/data/CrowdHuman/val/Images/283647,17890002197d807.jpg',
    '/root/data/CrowdHuman/val/Images/283992,76510000376e4d6.jpg',
    '/root/data/CrowdHuman/val/Images/284193,ff25000b6a403e9.jpg',
]


def eval_all(args):
    # model_path
    saveDir = config.model_dir
    evalDir = config.eval_dir
    misc_utils.ensure_dir(evalDir)
    model_file = os.path.join(saveDir, 
            'dump-{}.pth'.format(args.resume_weights))
    assert os.path.exists(model_file)
    # get devices
    str_devices = args.devices
    devices = misc_utils.device_parser(str_devices)

    inference(model_file, devices[0], image_list, score_thr=args.score_thr)


def inference(model_file, device, image_list, score_thr=0.3):
    torch.set_default_tensor_type('torch.FloatTensor')
    net = network.Network()
    net.cuda(device)
    check_point = torch.load(model_file)
    net.load_state_dict(check_point['state_dict'])
    net.eval()
    for image_path in image_list:
        np.set_printoptions(precision=2, suppress=True)
        img = cv2.imread(image_path, cv2.IMREAD_COLOR)
        image, im_info = get_data(img, device)
        pred_boxes = net(image, im_info)
        if if_set_nms:
            from set_nms_utils import set_cpu_nms
            n = pred_boxes.shape[0] // 2
            idents = np.tile(np.arange(n)[:,None], (1, 2)).reshape(-1, 1)
            pred_boxes = np.hstack((pred_boxes, idents))
            keep = pred_boxes[:, -2] > 0.05
            pred_boxes = pred_boxes[keep]
            keep = set_cpu_nms(pred_boxes, 0.5)
            pred_boxes = pred_boxes[keep]
        else:
            import det_tools_cuda as dtc
            nms = dtc.nms
            keep = nms(pred_boxes[:, :4], pred_boxes[:, 4], 0.5)
            pred_boxes = pred_boxes[keep]
            pred_boxes = np.array(pred_boxes)
            keep = pred_boxes[:, -1] > 0.05
            pred_boxes = pred_boxes[keep]
        width = int(im_info[0, -1])
        height = int(im_info[0, -2])
        dtboxes = boxes_dump(pred_boxes, False)
        pred_boxes = pred_boxes.tolist()
        filter_pre_boxes = []
        for pred_box in pred_boxes:
            pred_box[0] = int(round(pred_box[0]))
            pred_box[1] = int(round(pred_box[1]))
            pred_box[2] = int(round(pred_box[2]))
            pred_box[3] = int(round(pred_box[3]))
            if pred_box[4] < score_thr:
                continue
            filter_pre_boxes.append(pred_box)
        img = draw_boxes(img, filter_pre_boxes)
        cv2.imshow('image', img)
        key = cv2.waitKey()
        if key == 27 or key == ord('q') or key == ord('Q'):
            break


def boxes_dump(boxes, is_gt):
    result = []
    boxes = boxes.tolist()
    for box in boxes:
        if is_gt:
            box_dict = {}
            box_dict['box'] = [box[0], box[1], box[2]-box[0], box[3]-box[1]]
            box_dict['tag'] = box[-1]
            result.append(box_dict)
        else:
            box_dict = {}
            box_dict['box'] = [box[0], box[1], box[2]-box[0], box[3]-box[1]]
            box_dict['tag'] = 1
            box_dict['proposal_num'] = box[-1]
            box_dict['score'] = box[-2]
            result.append(box_dict)
    return result


def get_data(image, device):
    if config.eval_resize == False:
        resized_img, scale = image, 1
    else:
        resized_img, scale = dataset.resize_img_by_short_and_max_size(
            image, config.eval_image_short_size, config.eval_image_max_size)

    original_height, original_width = image.shape[0:2]
    height, width = resized_img.shape[0:2]
    transposed_img = np.ascontiguousarray(
        resized_img.transpose(2, 0, 1)[None, :, :, :],
        dtype=np.float32)
    im_info = np.array([height, width, scale, original_height, original_width],
                       dtype=np.float32)
    image = torch.Tensor(transposed_img).cuda(device)
    im_info = torch.Tensor(im_info[None, :]).cuda(device)
    return image, im_info


def run_test():
    parser = argparse.ArgumentParser()
    parser.add_argument('--resume_weights', '-r', default=None, type=str)
    parser.add_argument('--devices', '-d', default='0', type=str)
    parser.add_argument('--score_thr', default=0.7, type=float)
    args = parser.parse_args()
    eval_all(args)


if __name__ == '__main__':
    run_test()

